# *-* coding:utf8 *-*

from PyQt5.QtWidgets import QWidget
from PyQt5.QtWidgets import QHBoxLayout
from PyQt5.QtWidgets import QVBoxLayout
from PyQt5.QtWidgets import QComboBox
from PyQt5.QtWidgets import QLabel
from PyQt5.QtWidgets import QListView
from PyQt5.QtWidgets import QTextEdit
from PyQt5.QtWidgets import QLineEdit
from PyQt5.QtWidgets import QPushButton
from PyQt5.QtWidgets import QCheckBox
from PyQt5.QtWidgets import QGroupBox
from PyQt5.QtWidgets import QMessageBox
from PyQt5.QtWidgets import QAbstractItemView
from PyQt5.QtCore import QStringListModel
from PHFridaTools import FridaTools
from PyQt5.Qt import QCompleter
from PyQt5.Qt import Qt
import json
import traceback


class AndroidException:
    @staticmethod
    def format_exception_msg(widget, msg):
        QMessageBox.critical(widget, "错误", str(msg), QMessageBox.Yes | QMessageBox.No, QMessageBox.Yes)
        return 'traceback.format_exc():\n%s' % traceback.format_exc()


class AndroidTools(QWidget):

    def __init__(self, log_out):
        super(QWidget, self).__init__()
        self.log_out = log_out

    def insert_exception(self, exception):
        self.log_out.append(exception)


class AndroidFileOption(AndroidTools):

    def __init__(self, log_out):
        super(QWidget, self).__init__()
        self.log_out = log_out
        self.init_ui()

    def init_ui(self):
        pass


class AndroidRPC(AndroidTools):

    def __init__(self, log_out):
        super(AndroidRPC, self).__init__(log_out)
        self.log_out = log_out

        v_box = QVBoxLayout()
        call = QPushButton('test')
        call.clicked.connect(self.test)
        v_box.addWidget(call)

        h_temp = QHBoxLayout()
        temp_label = QLabel('参数1：')
        param1 = QLineEdit()
        h_temp.addWidget(temp_label)
        h_temp.addWidget(param1)
        v_box.addLayout(h_temp)

        h_temp = QHBoxLayout()
        temp_label = QLabel('参数2：')
        param2 = QLineEdit()
        h_temp.addWidget(temp_label)
        h_temp.addWidget(param2)
        v_box.addLayout(h_temp)

        h_temp = QHBoxLayout()
        temp_label = QLabel('参数3：')
        param3 = QLineEdit()
        h_temp.addWidget(temp_label)
        h_temp.addWidget(param3)
        v_box.addLayout(h_temp)

        h_temp = QHBoxLayout()
        temp_label = QLabel('参数4：')
        param4 = QLineEdit()
        h_temp.addWidget(temp_label)
        h_temp.addWidget(param4)
        v_box.addLayout(h_temp)

        self.setLayout(v_box)

    def test(self):
        pass

    def on_message(self, message, data):
        pass


class AndroidCapture(AndroidTools):

    def __init__(self, log_out):
        super(AndroidCapture, self).__init__(log_out)
        self.method_list = list()

        self.process_combo = QComboBox()
        self.process_combo.addItem('进程名')
        self.process_combo.addItem('进程号')
        self.process_combo.currentTextChanged.connect(self.process_change)
        self.process_text = QLineEdit()
        self.process_text.setText('com.hx.hxsocket')

        self.hook_combo = QComboBox()
        self.hook_combo.addItem('附加')
        self.hook_combo.addItem('自启动')
        self.hook_combo.currentTextChanged.connect(self.connect_change)

        connect = QLabel('Hook方式:')
        self.connect_combo = QComboBox()
        self.connect_combo.addItem('usb')
        self.connect_combo.addItem('wifi')

        h_box = QHBoxLayout()
        h_box.addWidget(self.hook_combo)
        h_box.addWidget(self.process_combo)
        h_box.addWidget(self.process_text)
        h_box.addWidget(connect)
        h_box.addWidget(self.connect_combo)

        method_button = QPushButton('开始抓包')
        method_button.clicked.connect(self.capture)

        hook_button = QPushButton('发送数据')
        hook_button.clicked.connect(self.send)

        h_two_box = QHBoxLayout()
        h_two_box.addWidget(method_button)
        h_two_box.addWidget(hook_button)

        self.resp_view = QTextEdit()
        self.resp_view.setFontFamily('Consolas')
        self.resp_view.setFontPointSize(13)
        temp_box = QHBoxLayout()
        temp_box.addWidget(self.resp_view)

        temp_group_box = QGroupBox('请求')
        temp_group_box.setLayout(temp_box)

        h_list_box = QHBoxLayout()
        h_list_box.addWidget(temp_group_box)

        self.resp_view = QTextEdit()
        self.resp_view.setFontFamily('Consolas')
        self.resp_view.setFontPointSize(13)
        temp_box = QHBoxLayout()
        temp_box.addWidget(self.resp_view)

        temp_group_box = QGroupBox('响应')
        temp_group_box.setLayout(temp_box)

        h_text_box = QHBoxLayout()
        h_text_box.addWidget(temp_group_box)

        v_box = QVBoxLayout()
        v_box.addLayout(h_box)
        v_box.addLayout(h_two_box)
        v_box.addLayout(h_list_box)
        v_box.addLayout(h_text_box)
        self.setLayout(v_box)

    def capture(self):
        pass

    def send(self):
        pass

    def process_change(self):
        pass

    def connect_change(self):
        pass

    def on_message(self, message, data):
        pass

    def on_message(self, message, data):
        if message['type'] == 'send':
            text = message['payload']
            js = json.loads(text)
            if js['name'] == 'functions':
                self.view.update_list_view(js)
            elif js['name'] == 'hook':
                self.view.update_text_view(js)
        else:
            print(message)


class FridaOption(QWidget):

    def __init__(self, log_out):
        super(QWidget, self).__init__()
        self.log_out = log_out

        self.start = None
        self.pause = None
        self.stop = None

        self.process_combo = QComboBox()
        self.process_combo.addItem('进程名')
        self.process_combo.addItem('进程号')
        self.process_combo.currentTextChanged.connect(self.process_change)
        self.process_text = QLineEdit()
        self.process_text.setText('com.hx.hxsocket')

        completer = QCompleter()
        string_list = QStringListModel()
        string_list.setStringList(['md', 'monb', 'mobileqq', 'dsfsd', '432'])

        completer.setModel(string_list)
        completer.setCaseSensitivity(Qt.CaseInsensitive)
        completer.setFilterMode(Qt.MatchRecursive)
        completer.setCompletionMode(QCompleter.UnfilteredPopupCompletion)
        self.process_text.setCompleter(completer)
        self.process_text.setToolTip('fds')

        self.hook_combo = QComboBox()
        self.hook_combo.addItem('附加')
        self.hook_combo.addItem('自启动')
        self.hook_combo.currentTextChanged.connect(self.hook_change)

        self.connect_combo = QComboBox()
        self.connect_combo.addItem('usb')
        self.connect_combo.addItem('wifi')

        h_box = QHBoxLayout()
        h_box.addWidget(self.hook_combo)
        h_box.addWidget(self.process_combo)
        h_box.addWidget(self.process_text)
        h_box.addWidget(self.connect_combo)

        start = QPushButton('开始')
        pause = QPushButton('暂停')
        stop = QPushButton('停止')
        h_box.addWidget(start)
        h_box.addWidget(pause)
        h_box.addWidget(stop)

        self.setLayout(h_box)

    def process_change(self):
        pass

    def hook_change(self):
        pass

    def register_start(self, callback):
        self.start = callback

    def register_pause(self, callback):
        self.pause = callback

    def register_stop(self, callback):
        self.stop = callback

    def start(self):
        self.start()

    def pause(self):
        self.pause()

    def stop(self):
        self.stop()


class AndroidRemoteHook(AndroidTools):

    def __init__(self, log_out):
        super(AndroidRemoteHook, self).__init__(log_out)
        self.method_list = list()
        self.method_mapping = dict()
        self.is_pause = True

        self.process_combo = QComboBox()
        self.process_combo.addItem('进程名')
        self.process_combo.addItem('进程号')
        self.process_combo.currentTextChanged.connect(self.process_change)
        self.process_text = QLineEdit()
        self.process_text.setText('com.hx.hxsocket')

        self.hook_combo = QComboBox()
        self.hook_combo.addItem('附加')
        self.hook_combo.addItem('自启动')
        self.hook_combo.currentTextChanged.connect(self.hook_change)

        connect = QLabel('Hook方式:')
        self.connect_combo = QComboBox()
        self.connect_combo.addItem('usb')
        self.connect_combo.addItem('wifi')

        h_box = QHBoxLayout()
        h_box.addWidget(self.hook_combo)
        h_box.addWidget(self.process_combo)
        h_box.addWidget(self.process_text)
        h_box.addWidget(connect)
        h_box.addWidget(self.connect_combo)

        self.start_button = QPushButton('开始')
        self.start_button.clicked.connect(self.start)
        self.pause_button = QPushButton('暂停')
        self.pause_button.setEnabled(False)
        self.pause_button.clicked.connect(self.pause)
        self.stop_button = QPushButton('关闭')
        self.stop_button.setEnabled(False)
        self.stop_button.clicked.connect(self.stop)
        h_box.addWidget(self.start_button)
        h_box.addWidget(self.pause_button)
        h_box.addWidget(self.stop_button)

        class_label = QLabel('请输入类名:')
        self.class_line_text = QLineEdit()
        self.class_line_text.setText('com.hx.hxsocket.MainActivity')

        self.batch = QCheckBox('批量挂钩')
        self.batch.stateChanged.connect(self.batch_click)
        self.overload = QCheckBox('重载函数挂钩')
        self.overload.stateChanged.connect(self.overload_click)

        self.method_button = QPushButton('获取类方法')
        self.method_button.clicked.connect(self.methods)

        self.hook_button = QPushButton('方法挂钩')
        self. hook_button.clicked.connect(self.hook)

        h_two_box = QHBoxLayout()
        h_two_box.addWidget(class_label)
        h_two_box.addWidget(self.class_line_text)
        h_two_box.addWidget(self.batch)
        h_two_box.addWidget(self.overload)
        h_two_box.addWidget(self.method_button)
        h_two_box.addWidget(self.hook_button)

        self.list_view = QListView()
        self.list_view.setEditTriggers(QAbstractItemView.NoEditTriggers)

        temp_box = QHBoxLayout()
        temp_box.addWidget(self.list_view)

        temp_group_box = QGroupBox('函数列表')
        temp_group_box.setLayout(temp_box)

        h_list_box = QHBoxLayout()
        h_list_box.addWidget(temp_group_box)

        self.text_view = QTextEdit()
        self.text_view.setFontFamily('Consolas')
        self.text_view.setFontPointSize(12)
        temp_box = QHBoxLayout()
        temp_box.addWidget(self.text_view)

        temp_group_box = QGroupBox('输出结果')
        temp_group_box.setLayout(temp_box)

        h_text_box = QHBoxLayout()
        h_text_box.addWidget(temp_group_box)

        v_box = QVBoxLayout()
        v_box.addLayout(h_box)
        v_box.addLayout(h_two_box)
        v_box.addLayout(h_list_box)
        v_box.addLayout(h_text_box)
        self.setLayout(v_box)

        self.hook_tools = FridaTools()
        self.hook_tools.set_script_path('script/frida/hook_script.js')
        self.hook_tools.load_script()
        self.hook_tools.devices(True, '')

    def batch_click(self):
        if self.batch.isChecked():
            self.list_view.setSelectionMode(QAbstractItemView.MultiSelection)
            self.overload.setEnabled(False)
        else:
            self.list_view.setSelectionMode(QAbstractItemView.SingleSelection)
            self.overload.setEnabled(True)

    def overload_click(self):
        if self.overload.isChecked():
            self.batch.setEnabled(False)
        else:
            self.batch.setEnabled(True)

    def update_list_view(self, data):
        if data['data'] == 'null':
            return
        data = json.loads(data['data'])
        self.method_list.clear()
        test = []
        prev = ""
        temp = ""
        for item in data:
            current = self.analysis_method_name(item)
            if current in prev:
                index = int(prev.split('_')[1]) + 1
                temp = current + '_' + str(index)
                test.append(temp)
            else:
                temp = current + '_0'
                test.append(temp)

            prev = temp
            self.method_mapping.update({item: temp})
            self.method_list.append(item)

        list_model = QStringListModel()
        list_model.setStringList(self.method_list)
        self.list_view.setModel(list_model)

    def update_text_view(self, data):
        try:
            self.text_view.append('函数名: ' + data['method'])
            arguments = json.loads(data['arguments'])
            self.text_view.append('参数 : ')
            for item in arguments:
                self.text_view.append('    ' + item)

            if 'return' in data.keys():
                self.text_view.append('返回值: ' + data['return'])
        except Exception as e:
            self.insert_exception(AndroidException.format_exception_msg(self, e))

    def methods(self):
        class_name = self.class_line_text.text()
        try:
            if class_name == '':
                return
            self.rpc_get_class_methods(class_name)
        except Exception as e:
            self.insert_exception(AndroidException.format_exception_msg(self, e))

    def hook(self):
        try:
            if not self.hook_tools:
                return

            if len(self.method_list) <= 0:
                return

            methods = dict()
            methods.update({'class_name': self.class_line_text.text()})
            if self.overload.isChecked():
                if len(self.list_view.selectedIndexes()) > 1:
                    raise Exception('多选bug')
                data = self.list_view.currentIndex().data()
                current = self.analysis_method_name(data)

                count = 0
                for item in self.method_list:
                    if current == self.analysis_method_name(item):
                        count = count + 1
                count = count - 1

                methods.update({'cmd': 'overload'})
                methods.update({'method': current})
                methods.update({'count': count})
                print(methods)
            elif self.batch.isChecked():
                methods.update({'cmd': 'batch'})
                temp = []
                for item in self.list_view.selectedIndexes():
                    temp.append(self.method_mapping[item.data()])
                methods.update({'methods': temp})

            else:
                if len(self.list_view.selectedIndexes()) > 1:
                    raise Exception('多选bug')
                method = self.method_mapping[self.list_view.currentIndex().data()]
                methods.update({'cmd': 'general'})
                methods.update({'method': method.split("_")[0]})
                methods.update({'index': method.split("_")[1]})
            self.rpc_get_hook_methods(str(methods))
        except Exception as e:
            self.insert_exception(AndroidException.format_exception_msg(self, e))

    def start(self):
        text = self.process_text.text()
        if text == '':
            return
        try:
            if self.hook_combo.currentText() == '自启动':
                self.hook_tools.swap_app(text, self.on_message)
            else:
                if self.process_combo.currentText() == '进程名':
                    self.hook_tools.attach_app(text, self.on_message)
                else:
                    self.hook_tools.attach_app(int(text), self.on_message)
            self.hook_button.setEnabled(True)
            self.method_button.setEnabled(True)
            self.pause_button.setEnabled(True)
            self.stop_button.setEnabled(True)
        except Exception as e:
            self.insert_exception(AndroidException.format_exception_msg(self, e))

    def pause(self):
        try:
            if self.is_pause:
                self.hook_tools.unload()
                self.pause_button.setText('启动')
                self.is_pause = False
            else:
                self.hook_tools.load(self.on_message)
                self.pause_button.setText('暂停')
                self.is_pause = True
        except Exception as e:
            self.insert_exception(AndroidException.format_exception_msg(self, e))

    def stop(self):
        try:
            self.hook_tools.detach()
            self.hook_button.setEnabled(False)
            self.method_button.setEnabled(False)
            self.pause_button.setEnabled(False)
            self.stop_button.setEnabled(False)
        except Exception as e:
            self.insert_exception(AndroidException.format_exception_msg(self, e))

    def rpc_get_class_methods(self, class_name):
            self.hook_tools.get_exports().get_function(class_name)

    def rpc_get_hook_methods(self, methods):
            self.hook_tools.get_exports().hook_method(methods)

    def process_change(self, text):
        pass

    def hook_change(self, text):
        if text == '自启动':
            self.process_combo.setCurrentText('进程名')
            self.process_combo.setEnabled(False)
        else:
            self.process_combo.setEnabled(True)

    def get_overload_methods(self, text):
        methods = []
        for method in self.method_list:
            if text in method:
                methods.append(method)
        return methods

    def on_message(self, message, data):
        if message['type'] == 'send':
            text = message['payload']
            js = json.loads(text)
            if js['name'] == 'functions':
                self.update_list_view(js)
            elif js['name'] == 'hook':
                self.update_text_view(js)
        else:
            print(message)

    @staticmethod
    def analysis_method_name(text):
        text = text.split('(')
        text = text[0]
        start = text.rfind('.')
        return text[start + 1:]
